# !/usr/bin/env python

import math
from os import P_NOWAIT

from matplotlib.cbook import pts_to_midstep

from easy_markers.generator import Marker, MarkerGenerator

from geometry_msgs.msg import Point, Vector3


# from kalman_filter import Kalman

from people_msgs.msg import People, Person, PositionMeasurementArray

import rospy


def distance(leg1, leg2):
    return math.sqrt(math.pow(leg1.x - leg2.x, 2) +
                     math.pow(leg1.y - leg2.y, 2) +
                     math.pow(leg1.z - leg2.z, 2))


def average(leg1, leg2):
    return Point((leg1.x + leg2.x) / 2,
                 (leg1.y + leg2.y) / 2,
                 (leg1.z + leg2.z) / 2)


def add(v1, v2):
    return Vector3(v1.x + v2.x, v1.y + v2.y, v1.z + v2.z)


def subtract(v1, v2):
    return Vector3(v1.x - v2.x, v1.y - v2.y, v1.z - v2.z)


def scale(v, s):
    v.x *= s
    v.y *= s
    v.z *= s


def printv(v):
    print('%.2f %.2f %.2f' % (v.x, v.y, v.z),)


gen = MarkerGenerator()
gen.type = Marker.ARROW
gen.ns = 'velocities'
gen.lifetime = .5

class Kalman:
    def __init__(self, Q=.002, R=1.0, P=.01):
        self.x = None
        self.Q = Q
        self.R = R
        self.P = P
        self.p = None

    def update(self, values):
        N = len(values)
        if self.x is None:
            self.x = values
            self.p = [self.P] * N
        else:
            for i in range(N):
                self.p[i] += self.Q
                k = self.p[i] / (self.p[i] + self.R)
                self.x[i] += k * (values[i] - self.x[i])
                self.p[i] = (1 - k) * self.p[i]

    def values(self):
        return self.x

class PersonEstimate(object):
    def __init__(self, msg):
        self.pos = msg
        self.reliability = 0.1
        self.k = Kalman()

    def update(self, msg):
        last = self.pos
        self.pos = msg
        self.reliability = max(self.reliability, msg.reliability)
        # print(msg.reliability)
        # print("is rebiligty")
        ivel = subtract(self.pos.pos, last.pos)
        time = (self.pos.header.stamp - last.header.stamp).to_sec()
        scale(ivel, 1.0 / time)
        # print("karlman")
        self.k.update([ivel.x, ivel.y, ivel.z])

    def age(self):
        return self.pos.header.stamp

    def get_id(self):
        return self.pos.object_id

    def velocity(self):
        k = self.k.values()
        if k is None:
            return Vector3()
        v = Vector3(k[0], k[1], k[2])
        # print(v)
        # print("v")
        return v

    def publish_markers(self, pub):
        gen.scale = [.1, .3, 0]
        gen.color = [1, 1, 1, 1]
        vel = self.velocity()
        m = gen.marker(points=[self.pos.pos, add(self.pos.pos, vel)])
        m.header = self.pos.header
        pub.publish(m)

    def get_person(self):
        p = Person()
        p.name = self.get_id()
        p.position = self.pos.pos
        p.velocity = self.velocity()
        p.reliability = self.reliability
        return self.pos.header.frame_id, p


class VelocityTracker(object):
    def __init__(self):
        self.people = {}
        self.TIMEOUT = rospy.Duration(rospy.get_param('~timeout', 2.0))
        self.sub = rospy.Subscriber('/people_mean_Array',
                                    PositionMeasurementArray,
                                    self.pm_cb)
        self.mpub = rospy.Publisher('/visualization_marker',
                                    Marker,
                                    queue_size=10)
        self.ppub = rospy.Publisher('/people',
                                    People,
                                    queue_size=10)


    def pm_cb(self, msg):
        for pm in msg.people:

            if pm.object_id in self.people:
                self.people[pm.object_id].update(pm)
                # print("update")
            else:
                # print("else")
                p = PersonEstimate(pm)
                self.people[pm.object_id] = p
                # print("hh else")

    def spin(self):
        rate = rospy.Rate(10)
        while not rospy.is_shutdown():
            # Remove People Older Than timeout param
            now = rospy.Time.now()
            for p in list(self.people.values()):
                if now - p.age() > self.TIMEOUT:
                    del self.people[p.get_id()]
            self.publish()
            rate.sleep()

    def publish(self):
        gen.counter = 0
        pl = People()
        pl.header.frame_id = None
        for p in self.people.values():
            p.publish_markers(self.mpub)
            # print("here")
            frame, person = p.get_person()
            pl.header.frame_id = frame
            pl.people.append(person)
        print(pl)
        print("pl here")
        self.ppub.publish(pl)


if __name__ == '__main__':
    rospy.init_node('people_velocity_tracker')
    vt = VelocityTracker()
    vt.spin()
